#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File        : tngs_split_gene.py
@Author      : Bing Liang
@Email       : believer19940901@gmail.com
@Date        : 2025/11/17
@Description : 使用 GFF 文件从基因组 FASTA 中拆分各基因序列并输出为 FASTA 文件
"""

import gzip
import re
from argparse import ArgumentParser, Namespace
from pathlib import Path
from collections import defaultdict

import pandas as pd
from Bio import SeqIO
from Bio.Seq import Seq


# ---------------------------------------------------------------
# 基因名转义（用于文件名，不改变 FASTA header）
# ---------------------------------------------------------------
def sanitize_gene_name(name: str) -> str:
    """
    将基因名转换为文件名可用的安全形式。
    保留字母、数字、下划线、点、横线，其他字符全部转为 "_"
    """
    return re.sub(r"[^A-Za-z0-9_.-]", "_", name)


# ---------------------------------------------------------------
# 命令行参数解析
# ---------------------------------------------------------------
def _parse_args() -> Namespace:
    """解析命令行参数。"""
    parser = ArgumentParser(description="根据 GFF 拆分 genome 中的 gene 序列。")
    parser.add_argument("-i", "--in_dir", required=True, type=str,
                        help="包含 *.fna.gz 与 *.gff.gz 的输入目录，每个样本一个文件夹")
    parser.add_argument("-o", "--out_dir", required=True, type=str,
                        help="输出目录，基因序列将写入其中")
    return parser.parse_args()


# ---------------------------------------------------------------
# GFF 文件解析
# ---------------------------------------------------------------
def _parse_gff(gff_path: Path) -> dict:
    """解析 GFF 文件，提取 gene 特征。"""

    with gzip.open(gff_path, "rt") as handle:
        df = pd.read_csv(
            handle,
            sep="\t",
            comment="#",
            header=None,
            usecols=[0, 2, 3, 4, 6, 8],
            names=["acc", "gentype", "start", "end", "strand", "info"]
        )

    acc_genes = defaultdict(list)

    for _, row in df.iterrows():
        if row["gentype"] != "gene":
            continue

        # 提取 gene 名（兼容 gene=xxx; gene="xxx" 等）
        m = re.search(r"gene=([^;]+)", row["info"])
        if not m:
            continue

        gene_name = m.group(1).strip().strip('"\'')   # 彻底去除引号

        gene_info = {
            "acc": row["acc"],
            "gentype": row["gentype"],
            "start": int(row["start"]),
            "end": int(row["end"]),
            "strand": row["strand"],
            "gene": gene_name,
            "info": row["info"],
        }

        acc_genes[row["acc"]].append(gene_info)

    return acc_genes


# ---------------------------------------------------------------
# 写出基因序列
# ---------------------------------------------------------------
def _write_gene_sequences(accs: dict, fna_path: Path, out_dir: Path) -> None:
    """从基因组序列中提取 gene 序列并写入文件。"""

    with gzip.open(fna_path, "rt") as handle:
        for record in SeqIO.parse(handle, "fasta"):
            chrom = record.id
            gene_list = accs.get(chrom)
            if not gene_list:
                continue

            for g in gene_list:

                # GFF 为 1-based，Biopython 为 0-based
                start = g["start"] - 1
                end = g["end"]

                seq: Seq = record.seq[start:end]

                if g["strand"] == "-":
                    seq = seq.reverse_complement()

                # ----------- 核心：基因名转义 -----------
                safe_gene = sanitize_gene_name(g["gene"])
                out_fp = out_dir / f"{safe_gene}.fas"
                # ----------------------------------------

                # header 中保留原基因名，不替换
                with open(out_fp, "a") as w:
                    w.write(f">{g['acc']} | {g['info']}\n{seq}\n")


# ---------------------------------------------------------------
# 主流程
# ---------------------------------------------------------------
def main() -> None:
    args = _parse_args()

    in_dir = Path(args.in_dir).absolute()
    out_dir = Path(args.out_dir).absolute()
    out_dir.mkdir(parents=True, exist_ok=True)

    for sample_dir in in_dir.iterdir():
        if not sample_dir.is_dir():
            continue

        fna_list = list(sample_dir.glob("*.fna.gz"))
        gff_list = list(sample_dir.glob("*.gff.gz"))

        if not fna_list or not gff_list:
            continue

        fna = fna_list[0]
        gff = gff_list[0]

        print(f"[INFO] 处理样本目录: {sample_dir.name}")
        print(f"[INFO] 使用基因组文件: {fna.name}")
        print(f"[INFO] 使用 GFF 文件: {gff.name}")

        accs = _parse_gff(gff)
        _write_gene_sequences(accs, fna, out_dir)

    print("[INFO] 处理完成。")


if __name__ == "__main__":
    main()
